package cn.feizhu.aimarket.controller;

import cn.dev33.satoken.annotation.SaCheckLogin;
import cn.dev33.satoken.annotation.SaIgnore;
import cn.dev33.satoken.stp.StpUtil;
import cn.feizhu.aimarket.ai.client.WriterChatClientFactory;
import cn.feizhu.aimarket.ai.workflow.config.process.WorkflowGraphProcess;
import cn.feizhu.aimarket.ai.workflow.model.constant.WorkflowControlConstant;
import cn.feizhu.aimarket.common.BaseResponse;
import cn.feizhu.aimarket.common.ResultUtils;
import cn.feizhu.aimarket.exception.ErrorCode;
import cn.feizhu.aimarket.exception.ThrowUtils;
import cn.feizhu.aimarket.model.dto.ai.GenerateFullBookRequest;
import cn.feizhu.aimarket.model.dto.ai.GeneratePagesRequest;
import cn.feizhu.aimarket.model.dto.ai.HumanFeedbackRequest;
import cn.feizhu.aimarket.ratelimit.annotation.RateLimit;
import cn.feizhu.aimarket.ratelimit.enums.RateLimitType;
import cn.hutool.core.util.StrUtil;
import com.alibaba.cloud.ai.graph.*;
import com.alibaba.cloud.ai.graph.async.AsyncGenerator;
import com.alibaba.cloud.ai.graph.checkpoint.config.SaverConfig;
import com.alibaba.cloud.ai.graph.checkpoint.constant.SaverConstant;
import com.alibaba.cloud.ai.graph.checkpoint.savers.MemorySaver;
import com.alibaba.cloud.ai.graph.exception.GraphRunnerException;
import com.alibaba.cloud.ai.graph.exception.GraphStateException;
import com.alibaba.cloud.ai.graph.state.StateSnapshot;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Sinks;

import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.TimeUnit;

import static reactor.core.publisher.Sinks.many;

@Slf4j
@RestController
@RequestMapping("/ai")
public class AiController {


    private final WriterChatClientFactory writerChatClientFactory;

    private final CompiledGraph compiledGraph;

    @Autowired
    public AiController(@Qualifier("contentFlowGraph") StateGraph stateGraph, WriterChatClientFactory writerChatClientFactory) throws GraphStateException {
        this.writerChatClientFactory = writerChatClientFactory;
        SaverConfig saverConfig = SaverConfig.builder().register(SaverConstant.MEMORY, new MemorySaver()).build();
        // TODO 这里照抄的human controller的中断点，后续可以根据需求调整
        this.compiledGraph = stateGraph
                .compile(CompileConfig.builder().saverConfig(saverConfig).interruptBefore(WorkflowControlConstant.HUMAN_FEEDBACK).build());
    }

    @SaCheckLogin
    @RateLimit(key = "ai.workflow", limitType = RateLimitType.USER, rate = 1, rateInterval = 1, timeUnit = TimeUnit.DAYS, message = "对话调用已达今日上限", vipRateOverride = 10)
    @PostMapping("/workflow")
    public BaseResponse<?> workFlow(@RequestBody GenerateFullBookRequest generateFullBookRequest) {
        String userMessage = generateFullBookRequest.getUserMessage();
        ThrowUtils.throwIf(StrUtil.isBlank(userMessage), ErrorCode.PARAMS_ERROR, "用户消息不能为空");
        ThrowUtils.throwIf(userMessage.length() > 2000, ErrorCode.PARAMS_ERROR, "用户消息不能超过2000字");
        // TODO 目前使用userId作为threadId如何控制只生成一个，其实可以根据appId+userId生成一个threadId
        long userId = (long) StpUtil.getSession().get("userId");
        long appId = (long) StpUtil.getSession().get("appId");
        String threadId = userId + "_" + appId;
        RunnableConfig runnableConfig = RunnableConfig.builder().threadId(threadId).build();


        try {
            // 构建输入参数Map TODO 这里必须把data用到的东西都放进去初始化，否则后面添加没有
            Map<String, Object> inputMap = new HashMap<>();
            inputMap.put(WorkflowControlConstant.USER_INPUT, userMessage);
            // 这里保存应用ID和用户ID作为后续构建路径 TODO 全部用String格式
            inputMap.put(WorkflowControlConstant.APP_ID, String.valueOf(appId));
            inputMap.put(WorkflowControlConstant.USER_ID, String.valueOf(userId));
            // TODO 这里要进行初始化重置之前的workflow吗？
            // 调用编译后的图执行workflow
            Optional<OverAllState> result = this.compiledGraph.invoke(inputMap, runnableConfig);
            // 返回结果
            return ResultUtils.success(result.map(OverAllState::data).orElse(new java.util.HashMap<>()));
        } catch (Exception e) {
            ThrowUtils.throwIf(true, ErrorCode.SYSTEM_ERROR, "工作流执行失败: " + e.getMessage());
            return null;
        }
    }

    /**
     * 尝试写流式返回的
     * TODO EventSource只支持Get
     * @param generateFullBookRequest
     * @return
     */
    @SaCheckLogin
    @RateLimit(key = "ai.stream", limitType = RateLimitType.USER, rate = 1, rateInterval = 1, timeUnit = TimeUnit.DAYS, message = "对话调用已达今日上限", vipRateOverride = 10)
    @GetMapping("/streamflow")
    public Flux<ServerSentEvent<String>> workFlowStream(@RequestParam(value = "generateFullBookRequest") String generateFullBookRequest) throws GraphRunnerException {
        ThrowUtils.throwIf(StrUtil.isBlank(generateFullBookRequest), ErrorCode.PARAMS_ERROR, "用户消息不能为空");
        ThrowUtils.throwIf(generateFullBookRequest.length() > 2000, ErrorCode.PARAMS_ERROR, "用户消息不能超过2000字");
        // TODO 目前使用userId作为threadId如何控制只生成一个，其实可以根据appId+userId生成一个threadId
        long userId = (long) StpUtil.getSession().get("userId");
        long appId = (long) StpUtil.getSession().get("appId");
        String threadId = userId + "_" + appId;
        RunnableConfig runnableConfig = RunnableConfig.builder().threadId(threadId).build();
        WorkflowGraphProcess workflowGraphProcess = new WorkflowGraphProcess(compiledGraph);

        Map<String, Object> inputMap = new HashMap<>();
        inputMap.put(WorkflowControlConstant.USER_INPUT, generateFullBookRequest);
        // 这里保存应用ID和用户ID作为后续构建路径 TODO 全部用String格式
        inputMap.put(WorkflowControlConstant.APP_ID, String.valueOf(appId));
        inputMap.put(WorkflowControlConstant.USER_ID, String.valueOf(userId));
        AsyncGenerator<NodeOutput> asyncGenerator = this.compiledGraph.stream(inputMap, runnableConfig);
        Sinks.Many<ServerSentEvent<String>> sink = many().unicast().onBackpressureBuffer();
        workflowGraphProcess.processStream(asyncGenerator, sink);
        return sink.asFlux()
                .doOnCancel(() -> {
                    log.info("Client disconnected from stream");
                })
                .doOnError(e->{
                    log.error("Error occurred during streaming", e);
                });

    }


    @SaCheckLogin
    @GetMapping("/feedback")
    public Flux<ServerSentEvent<String>> feedback(@RequestParam(value = "feedback") String feedback) throws GraphRunnerException {

        ThrowUtils.throwIf(StrUtil.isBlank(feedback), ErrorCode.PARAMS_ERROR, "用户消息不能为空");
        ThrowUtils.throwIf(feedback.length() > 2000, ErrorCode.PARAMS_ERROR, "用户消息不能超过2000字");
        long userId = (long) StpUtil.getSession().get("userId");
        long appId = (long) StpUtil.getSession().get("appId");
        String threadId = userId + "_" + appId;
        RunnableConfig runnableConfig = RunnableConfig.builder().threadId(threadId).build();
        StateSnapshot stateSnapshot = this.compiledGraph.getState(runnableConfig);
        OverAllState state = stateSnapshot.state();
        state.withResume();
        Map<String, Object> inputMap = new HashMap<>();
        inputMap.put(WorkflowControlConstant.FEEDBACK, feedback);
        state.withHumanFeedback(new OverAllState.HumanFeedback(inputMap, ""));

        Sinks.Many<ServerSentEvent<String>> sinks = Sinks.many().unicast().onBackpressureBuffer();
        WorkflowGraphProcess workflowGraphProcess = new WorkflowGraphProcess(compiledGraph);
        AsyncGenerator<NodeOutput> asyncGenerator = compiledGraph.streamFromInitialNode(state, runnableConfig);
        workflowGraphProcess.processStream(asyncGenerator, sinks);
        return sinks.asFlux()
                .doOnCancel(() -> {
                    log.info("Client disconnected from stream");
                })
                .doOnError(e->{
                    log.error("Error occurred during streaming", e);
                });

    }

    /**
     * 废弃方法，无法配合流式输出返回，所以直接作为废弃路径
     * @param humanFeedbackRequest
     * @return
     */
    @Deprecated
    @SaCheckLogin
    @PostMapping("/feedback_OUTDATED")
    public BaseResponse<?> feedBack_OUTDATED(@RequestBody HumanFeedbackRequest humanFeedbackRequest) {
        /**
         * 如果用户同意大纲那么返回 APPROVED
         * 如果用户不同意大纲那么返回 NEEDS_REVISION
         * TODO 主要是这里想要测试一下智能判断节点
         */
        String feedback = humanFeedbackRequest.getUserMessage();
        ThrowUtils.throwIf(StrUtil.isBlank(feedback), ErrorCode.PARAMS_ERROR, "用户消息不能为空");
        ThrowUtils.throwIf(feedback.length() > 2000, ErrorCode.PARAMS_ERROR, "用户消息不能超过2000字");
        // TODO 这里跟上面创建一样的runnableConfig
        long userId = (long) StpUtil.getSession().get("userId");
        long appId = (long) StpUtil.getSession().get("appId");
        String threadId = userId + "_" + appId;
        RunnableConfig runnableConfig = RunnableConfig.builder().threadId(threadId).build();
        // TODO 这里是获取stateSnapShot 但是下面的resume方法已经包含了上面的操作了
//        StateSnapshot stateSnapshot = this.compiledGraph.getState(runnableConfig);
//        OverAllState state = stateSnapshot.state();
        // 这里设置为resume就是保留的意思
//        state.withResume();

        try {
            // 构建输入参数Map TODO 这里必须把data用到的东西都放进去初始化，否则后面添加没有
            Map<String, Object> inputMap = new HashMap<>();
            inputMap.put(WorkflowControlConstant.FEEDBACK, feedback);
            // 调用编译后的图执行workflow TODO 这里用state试一试 这里用resume试一试，这个源码看过
            Optional<OverAllState> result = this.compiledGraph.resume(new OverAllState.HumanFeedback(inputMap, ""), runnableConfig);
            // 返回结果
            return ResultUtils.success(result.map(OverAllState::data).orElse(new java.util.HashMap<>()));
        } catch (Exception e) {
            ThrowUtils.throwIf(true, ErrorCode.SYSTEM_ERROR, "工作流执行失败: " + e.getMessage());
            return null;
        }

    }

    @SaIgnore
    @GetMapping("/write")
    public BaseResponse<String> write(@RequestParam String userMessage) {
        System.out.println(userMessage);
        var chatClient = writerChatClientFactory.createWriterChatClient();
        ChatClient.CallResponseSpec call = chatClient.prompt()
                .user(userMessage)
                .call();
        System.out.println(call);
        String string = call.chatResponse().getResult().getOutput().toString();
        System.out.println(string);
        return ResultUtils.success(string);
    }


    @Resource
    private cn.feizhu.aimarket.ai.service.AiService aiService;

    /**
     * 并行生成多个文本，TODO 注意目前测试阶段，service只返回OK代表正常
     *
     * @param generatePagesRequest
     * @return
     */
    @SaCheckLogin
    @RateLimit(key = "ai.generatePages", limitType = RateLimitType.USER, rate = 1, rateInterval = 1, timeUnit = TimeUnit.DAYS, message = "生成次数已达今日上限", vipRateOverride = 10)
    @PostMapping("/generate/pages")
    public BaseResponse<String> generatePages(@RequestBody GeneratePagesRequest generatePagesRequest) {
        // 参数校验
        ThrowUtils.throwIf(generatePagesRequest == null, ErrorCode.PARAMS_ERROR, "请求参数不能为空");
        String userMessage = generatePagesRequest.getUserMessage();
        Integer parallelNum = generatePagesRequest.getParallelNum();

        ThrowUtils.throwIf(StrUtil.isBlank(userMessage), ErrorCode.PARAMS_ERROR, "用户消息不能为空");
        ThrowUtils.throwIf(parallelNum == null || parallelNum <= 0, ErrorCode.PARAMS_ERROR, "并行数量必须大于0");
        ThrowUtils.throwIf(parallelNum > 10, ErrorCode.PARAMS_ERROR, "并行数量不能超过10");

        /**
         * 通过StpSession获取用户ID和应用ID
         */
        long userId = (long) StpUtil.getSession().get("userId");
        long appId = (long) StpUtil.getSession().get("appId");
        String result = aiService.generatePages(userId, appId, userMessage, parallelNum);
        return ResultUtils.success(result);
    }


}
